#!/usr/bin/env python


from scapy.all import Dot11, rdpcap, sniff
import binascii
import itertools

class PacketDecoder:
    # C = 48
    C = 56
    def __init__(self, pkt):
        self.pkt = pkt

    @property
    def value(self):
        # return len(pkt[Dot11]["Dot11WEP"].wepdata)-PacketDecoder.C
        if not self.pkt[Dot11].haslayer("Raw"):
            return 0
        return len(self.pkt[Dot11]["Raw"].load)-PacketDecoder.C

    @property
    def sourceAddr(self):
        return self.pkt[Dot11].addr3

    @property
    def destAddr(self):
        return self.pkt[Dot11].addr1

    @property
    def transAddr(self):
        return self.pkt[Dot11].addr2

    @property
    def type(self):
        return self.pkt[Dot11].type
    


class EasyconfigStartFlagStateMachine:
    def __init__(self, decoder):
        self.state = "init"
        self.startedToAnalyze = False
        self.startedToCount525= False
        self.cnt = {5:0, 25:0}
        self.decoder=decoder

    def onPkt(self, pkt):
        # print pkt.value
        if self.state == "init":
            if pkt.value == 5 or pkt.value ==25:
                self.cnt[pkt.value] += 1
                self.state = "count525"

        if self.state == "count525":
            if pkt.value == 5 or pkt.value ==25:
                self.cnt[pkt.value] += 1
                if self.cnt[5]>=2 and self.cnt[25] >=2:
                    self.state = "analyze"
            else:
                self.state = "init"
                self.cnt = {5:0, 25:0}

        if self.state == "analyze":
            if pkt.value > 50:
                self.decoder.decodeSingleData(pkt.value)
                if self.decoder.isFinished():
                    self.state = "finished"

    def isFinished(self):
        return self.state == "finished"
            

class EasyconfigDecoder:
    def __init__(self):
        self.pos_got = set()
        self.wifi_info_bin = [0] * 1000
        self.orderd_origin = [0] * 1000
        self.state = "analyzing"
        self.wifo_info = ""


    def decodeSingleData(self, data):
        pos = ((data-50) & 0xfff8) / 8
        if pos<0 :
        # TODO:change this to throw exception
            print data
        self.pos_got.add(pos)

        self.orderd_origin[pos] = data
        val = (data-50) & 0x7
        for i in range(0,3):
            if val & (1<<i):
                self.wifi_info_bin[pos*3+i] = 1

        if len(self.pos_got) == max(self.pos_got) +1:
            self.checkResult()


    def isFinished(self):
        return self.state == "finished"

    def checkResult(self):
        self.wifi_info_bin = ''.join(str(c) for c in self.wifi_info_bin)
        wifi_info = ""

        for i in range(0,len(self.wifi_info_bin)/8):
            wifi_info += chr(int(self.wifi_info_bin[i*8:i*8+8][::-1],2))
        wifi_info_len = ((max(self.pos_got)-1)*3 +1) /8
        wifi_info = wifi_info[:wifi_info_len]
        # print wifi_info
        if ord(wifi_info[0]) == len(wifi_info) +1:
            self.wifi_info = wifi_info
            self.state = "finished"            

    def printResult(self):
        print repr(self.wifi_info)


dfa_dispatcher = {}

def dispatcher(pkt):
    # pkt.show()
    if not pkt.haslayer(Dot11):
        return 

    p = PacketDecoder(pkt)
    # if not p.type == 0: # must be data
        # return 
    if not p.destAddr == "01:00:5e:00:00:fb": # dest must be ipv4 multicast
        return 

    source_addr = p.sourceAddr
    if source_addr in dfa_dispatcher:
        dfa = dfa_dispatcher[source_addr]
        dfa.onPkt(p)
        if dfa.isFinished():
            dfa.decoder.printResult()
            exit(0)
    else:
        decoder = EasyconfigDecoder()
        dfa = EasyconfigStartFlagStateMachine(decoder)
        dfa_dispatcher[source_addr] = dfa
        dfa.onPkt(p)

    


def main():
    pkts = rdpcap("test.pcapng")
    # pkts = [p for p in pkts if p.haslayer(Dot11) and p[Dot11].addr3=="fc:64:ba:33:65:14"]
   
    for pkt in pkts:
        dispatcher(pkt)

    
    # print dfa_dispatcher["fc:64:ba:33:65:14"].decoder.checkResult()
    # print dfa_dispatcher
    # sniff(iface="en0", prn=dispatcher, count=100)


if __name__ =="__main__":
    main()